from PyQt5.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                           QPushButton, QFileDialog, QLabel, QListWidget)
from PyQt5.QtCore import Qt
from sql_parser import SQLParser
from er_generator import ERGenerator
from PyQt5.QtGui import QPainter, QPen, QBrush, QColor, QFont

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("SQL转ER图工具")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中心部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        
        # 创建左侧面板
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_panel.setMaximumWidth(250)
        main_layout.addWidget(left_panel)
        
        # 只保留导入和导出按钮
        self.import_btn = QPushButton("导入SQL文件")
        self.import_btn.clicked.connect(self.import_sql)
        left_layout.addWidget(self.import_btn)
        
        self.export_btn = QPushButton("导出ER图")
        self.export_btn.clicked.connect(self.export_er)
        self.export_btn.setEnabled(False)
        left_layout.addWidget(self.export_btn)
        
        # 添加表格列表
        self.table_list = QListWidget()
        self.table_list.itemClicked.connect(self.on_table_selected)
        left_layout.addWidget(QLabel("实体列表:"))
        left_layout.addWidget(self.table_list)
        
        # 存储表名和注释的映射关系
        self.table_map = {}
        self.tables = []
        
        # 文件路径显示标签
        self.file_label = QLabel("未选择文件")
        left_layout.addWidget(self.file_label)
        
        # 创建右侧绘图区域
        self.drawing_area = ERDrawingArea()
        main_layout.addWidget(self.drawing_area)
        
        # 设置布局比例
        main_layout.setStretch(0, 1)
        main_layout.setStretch(1, 4)
        
        self.er_generator = ERGenerator(self.drawing_area)
        self.drawing_area.setGenerator(self.er_generator)
        self.sql_content = None
    
    def import_sql(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择SQL文件", "", "SQL文件 (*.sql);;所有文件 (*.*)")
        
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    self.sql_content = f.read()
                self.file_label.setText(f"已导入: {file_path}")
                # 导入后直接生成ER图
                self.generate_er()
            except Exception as e:
                self.file_label.setText(f"文件读取错误: {str(e)}")
                self.sql_content = None
                self.export_btn.setEnabled(False)
    
    def generate_er(self):
        if not self.sql_content:
            return
            
        parser = SQLParser()
        self.tables = parser.parse(self.sql_content)
        
        if self.tables:
            # 更新表格列表
            self.table_list.clear()
            self.table_map.clear()
            
            for table in self.tables:
                display_name = table.comment or table.name
                self.table_map[display_name] = table
                self.table_list.addItem(display_name)
            
            # 生成第一个表的ER图
            self.er_generator.generate(self.tables)
            if self.tables:
                self.er_generator.show_table(self.tables[0].name)
            self.export_btn.setEnabled(True)
    
    def export_er(self):
        if self.er_generator:
            file_path, _ = QFileDialog.getSaveFileName(
                self, "导出ER图", "", "PNG文件 (*.png);;所有文件 (*.*)")
            if file_path:
                self.er_generator.export(file_path)
    
    def on_table_selected(self, item):
        display_name = item.text()
        # 通过映射获取表对象
        table = self.table_map.get(display_name)
        if table:
            print(f"切换到表: {table.name}")  # 调试信息
            self.er_generator.show_table(table.name)
            self.drawing_area.update()

# 修改 ERDrawingArea 类
class ERDrawingArea(QWidget):
    def __init__(self):
        super().__init__()
        self.er_generator = None
        self.setAutoFillBackground(True)
        palette = self.palette()
        palette.setColor(self.backgroundRole(), QColor(255, 255, 255))
        self.setPalette(palette)
        self.setMinimumSize(800, 600)
        self.setCursor(Qt.OpenHandCursor)  # 设置手形光标
        
    def setGenerator(self, generator):
        self.er_generator = generator
    
    def paintEvent(self, event):
        if self.er_generator:
            painter = QPainter(self)
            self.er_generator.paint(painter)
    
    def mousePressEvent(self, event):
        if not self.er_generator:
            return
            
        if event.button() == Qt.LeftButton:
            # 直接处理拖动，不需要Ctrl键
            self.er_generator.handle_mouse_press(event)
        elif event.button() == Qt.RightButton:
            self.er_generator.handle_mouse_press(event)
    
    def mouseMoveEvent(self, event):
        if not self.er_generator:
            return
            
        if hasattr(self.er_generator, 'selected_shape') and \
           self.er_generator.selected_shape and \
           self.er_generator.selected_shape.is_dragging:
            self.er_generator.handle_mouse_move(event)
    
    def mouseReleaseEvent(self, event):
        if not self.er_generator:
            return
            
        if event.button() == Qt.LeftButton and \
           hasattr(self.er_generator, 'selected_shape') and \
           self.er_generator.selected_shape:
            self.er_generator.handle_mouse_release(event)
    
    def wheelEvent(self, event):
        if not self.er_generator:
            return
            
        if event.modifiers() == Qt.ControlModifier:
            # Ctrl+滚轮进行缩放
            if event.angleDelta().y() > 0:
                self.er_generator.zoom_in()
            else:
                self.er_generator.zoom_out()
            self.update()
            event.accept()
        else:
            super().wheelEvent(event)